from python_ai.DL.tensorflow.x6_convolution.my.x_pic_data import x_pic_data
import tensorflow.compat.v1 as tf
import numpy as np
import matplotlib.pyplot as plt

idx = 2
pic = np.array([x_pic_data[idx]])
ph_pic = tf.placeholder(tf.float32, [None, pic.shape[1], pic.shape[2], pic.shape[3]], 'ph_pic')
print(ph_pic)

knl_raw = np.array([
    [[1, 0, -1],
     [1, 0, -1],
     [1, 0, -1],],
    [[1, 1, 1],
     [0, 0, 0],
     [-1, -1, -1]],
    [[-1, 0, 1],
     [-1, 0, 1],
     [-1, 0, 1],],
    [[-1, -1, -1],
     [0, 0, 0],
     [1, 1, 1]],
], dtype=np.float32)
print(f'Kernel data shape: {knl_raw.shape}')
kh, kw, cin, cout = knl_raw.shape[1], knl_raw.shape[2], pic.shape[3], knl_raw.shape[0]
print(f'Target shape: {kh, kw, cin, cout}')
knl = np.zeros([kh, kw, cin, cout])
for c_idx, c in enumerate(knl_raw):
    for h_idx, h in enumerate(c):
        for w_idx, w in enumerate(h):
            knl[h_idx, w_idx, 0, c_idx] = w
filter1 = tf.Variable(knl, dtype=tf.float32, name='filter1')
conv = tf.nn.conv2d(ph_pic, filter1, strides=[1, 1, 1, 1], padding='VALID')

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    conv_v = sess.run(conv, feed_dict={ph_pic: pic})

    plt.figure(figsize=[18, 12])
    spr = 2  # subplot row
    spc = 3  # subplot column
    spn = 0
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.imshow(pic[0])
    for ch in range(cout):
        spn += 1
        if spn > spr * spc:
            break
        plt.subplot(spr, spc, spn)
        p = conv_v[0, ..., ch]
        plt.imshow(p)
